import os
import time
import matplotlib.pyplot as plt
import csv
from loguru import logger

from models.World import World

# 指定支持中文的字体
plt.rcParams['font.family'] = 'SimHei'

# 解决负号显示问题
plt.rcParams['axes.unicode_minus'] = False

# 模拟参数
WIDTH = 25
HEIGHT = 25
NUM_FOODS = 1000  # 食物数量
NUM_ORGANISMS = 50  # 生物数量
NUM_GENERATIONS = 50  # 生物进化代数
NUM_SIMULATIONS = 20  # 模拟组数

# 创建保存目录
if not os.path.exists('images'):
    os.makedirs('images')
if not os.path.exists('data'):
    os.makedirs('data')


for sim_id in range(1, NUM_SIMULATIONS + 1):
    logger.info(f"开始第 {sim_id} 组模拟...")
    world = World(WIDTH, HEIGHT, NUM_FOODS, NUM_ORGANISMS)

    # 初始化当前模拟数据记录
    generations = []
    population_counts = []
    max_population_counts = 0
    food_counts = []
    avg_vision_levels = []
    avg_speed_levels = []

    # 用于记录种群数量是否持续减少
    decreasing_population = False
    prev_population = len(world.organisms)

    for generation in range(NUM_GENERATIONS):
        generations.append(generation + 1)
        current_population = len(world.organisms)
        population_counts.append(current_population)
        if current_population > max_population_counts:
            max_population_counts = current_population
        food_counts.append(len(world.foods))

        if world.organisms:
            visions = [o.genes[0] for o in world.organisms]
            speeds = [o.genes[1] for o in world.organisms]
            avg_vision = sum(visions) / len(visions)
            avg_speed = sum(speeds) / len(speeds)
        else:
            avg_vision = 0
            avg_speed = 0

        avg_vision_levels.append(avg_vision)
        avg_speed_levels.append(avg_speed)

        # 检查种群数量是否持续减少
        if current_population < prev_population:
            if decreasing_population:
                logger.warning(f"警告：第 {sim_id} 组模拟，第 {generation + 1} 代，生物种群数量持续减少！")
                # 采取适当措施，例如增加食物数量
                additional_food = NUM_FOODS * 40 % 100
                world.add_food(additional_food)  # 假设 World 类有 add_food 方法
                logger.success(f"已增加 {additional_food} 个食物以改善生物生存环境。")
            else:
                decreasing_population = True
        else:
            decreasing_population = False

        prev_population = current_population

        world.update()

    # 输出此次模拟数据
    logger.info(f"Max population counts: {max_population_counts}")

    # 保存当前模拟数据
    csv_filename = f'data/simulation_{sim_id:02d}.csv'
    with open(csv_filename, 'w', newline='') as csvfile:
        fieldnames = ['Generation', 'Population Count', 'Food Count',
                      'Average Vision Level', 'Average Speed Level']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for i in range(len(generations)):
            writer.writerow({
                'Generation': generations[i],
                'Population Count': population_counts[i],
                'Food Count': food_counts[i],
                'Average Vision Level': avg_vision_levels[i],
                'Average Speed Level': avg_speed_levels[i]
            })

    # 生成当前模拟图表
    fig = plt.figure(figsize=(12, 8))
    plt.subplot(2, 2, 1)
    plt.plot(generations, population_counts, marker='o', linestyle='-')
    plt.title(f'第 {sim_id} 组模拟 - 种群数量变化')
    plt.xlabel('世代')
    plt.ylabel('个体数量')

    plt.subplot(2, 2, 2)
    plt.plot(generations, food_counts, marker='o', color='orange', linestyle='-')
    plt.title('食物数量变化')
    plt.xlabel('世代')
    plt.ylabel('食物数量')

    plt.subplot(2, 2, 3)
    plt.plot(generations, avg_vision_levels, marker='o', color='green', linestyle='-')
    plt.title('平均视野等级变化')
    plt.xlabel('世代')
    plt.ylabel('视野等级 (0-2)')

    plt.subplot(2, 2, 4)
    plt.plot(generations, avg_speed_levels, marker='o', color='red', linestyle='-')
    plt.title('平均速度等级变化')
    plt.xlabel('世代')
    plt.ylabel('速度等级 (0-1)')

    plt.tight_layout()
    img_filename = f'images/simulation_{sim_id:02d}_charts.png'
    plt.savefig(img_filename)
    plt.close(fig)

    logger.success(f"第 {sim_id} 组模拟完成，数据已保存")

logger.success(f"所有 {NUM_SIMULATIONS} 组模拟全部完成！")
